package io.github.zerobyteword.oneasynctool;

import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

public class WorkNode<T> {
    private static final WrapResult NIL = new WrapResult(null);
    private Workable<T> workable;
    private ExecutorService executorService;
    private CountDownLatch countDownLatch;
    private Long timeOut;
    private boolean dependHasTimeOut;
    private boolean dependNotEmpty;
    // 存放结果或异常
    private volatile AtomicReference<Object> result = new AtomicReference<>(null);
    private boolean setDefault;
    private T defaultValue;
    private volatile AtomicBoolean executed = new AtomicBoolean(false);
    private volatile Future<?> future;

    private List<WorkNode<?>> dependWorkNodes;
    private Set<WorkNode<?>> nextWorkNodes;

    private WorkNode() {

    }

    public static <T> Builder<T> builder() {
        return new Builder<T>();
    }

    public Long getTimeOut() {
        return timeOut;
    }

    public void configRunInfo(ExecutorService executorService, CountDownLatch countDownLatch) {
        this.executorService = executorService;
        this.countDownLatch = countDownLatch;
    }

    public List<WorkNode<?>> getDependWorkNodes() {
        return dependWorkNodes;
    }

    public Set<WorkNode<?>> getNextWorkNodes() {
        return nextWorkNodes;
    }

    public T getResult() {
        Object r = this.result.get();
        if (r == null) {
            throw new ResultNotFoundException();
        }
        if (r instanceof WrapResult) {
            WrapResult wrapResult = (WrapResult) r;
            if (wrapResult.throwable == null) {
                return null;
            } else {
                throw new ExceptionWrapper(wrapResult.throwable);
            }
        }
        return (T) r;
    }


    private void tryRunNextWork() {
        boolean hasParentThread = this.future != null;
        if (this.nextWorkNodes != null && this.nextWorkNodes.size() > 0) {
            WorkNode<?> runInParentThreadNode = null;
            for (WorkNode<?> workNode : this.nextWorkNodes) {
                if (workNode.tryRunAndCheckInThread(hasParentThread)) {
                    runInParentThreadNode = workNode;
                    hasParentThread = false;
                }
            }
            if (runInParentThreadNode != null) {
                runInParentThreadNode.future = future;
                runInParentThreadNode.doWork();
            }
        }
    }


    private boolean tryRunAndCheckInThread(boolean hasParentThread) {
        if (!this.dependNotEmpty) {
            if (hasParentThread) {
                // 尝试变更状态
                return changeStateExecuted();
            } else {
                // 提交新线程
                submitNewTask();
            }
        }
        boolean allFinished = true;
        boolean hasException = false;
        Throwable execution = null;
        for (WorkNode<?> workNode : this.dependWorkNodes) {
            Object r = workNode.result.get();
            if (r == null) {
                allFinished = false;
                continue;
            }
            if (r instanceof WrapResult && ((WrapResult) r).throwable != null) {
                hasException = true;
                execution = ((WrapResult) r).throwable;
                break;
            }
        }
        if (hasException) {
            // 取消其他节点
            cancelDependNodes();
            if (hasParentThread) {
                return failValueSet(execution);
            } else {
                // 提交新线程
                submitNewTask();
            }
            return false;
        }
        if (allFinished) {
            if (hasParentThread) {
                // 尝试变更状态
                return changeStateExecuted();
            } else {
                // 提交新线程
                submitNewTask();
            }
        }
        return false;
    }

    private void cancelDependNodes() {
        if (this.dependNotEmpty) {
            for (WorkNode<?> workNode : this.dependWorkNodes) {
                workNode.cancel();
            }
        }
    }

    public void cancel() {
        if (this.result.get() == null) {
            this.changeStateExecuted();
            this.completeThrowable(new CancellationException());
            Future<?> temp = this.future;
            if (temp != null && !temp.isDone() && !temp.isCancelled()) {
                temp.cancel(false);
            }
        }
    }


    private boolean hasExecuted() {
        return this.executed.get();
    }

    public void submitNewTask() {
        if (hasExecuted()) {
            return;
        }
        if (changeStateExecuted()) {
            try {
                // 提交线程池
                Future<?> future = executorService.submit(this::doWork);
                this.future = future;
                tryTriggerNextWorkPreCreate();
            } catch (Throwable e) {
                // 提交线程池失败或者超时
                failHandle(e);
            }
        }
    }

    private boolean changeStateExecuted() {
        return executed.compareAndSet(false, true);
    }

    private void tryTriggerNextWorkPreCreate() {
        if (this.timeOut != null && this.nextWorkNodes != null && this.nextWorkNodes.size() > 0) {
            for (WorkNode<?> workNode : this.nextWorkNodes) {
                workNode.tryPreSubmitTask();
            }
        }
    }

    private void tryPreSubmitTask() {
        if (this.dependHasTimeOut) {
            boolean hasFutureNotInit = this.dependWorkNodes
                    .stream().anyMatch(e -> e.future == null);
            if (!hasFutureNotInit) {
                submitNewTask();
            }
        }
    }

    final boolean completeValue(T value) {
        return this.result.compareAndSet(null, value == null ? NIL : value);
    }

    final boolean completeThrowable(Throwable x) {
        return this.result.compareAndSet(null, encodeThrowable(x));
    }

    private WrapResult encodeThrowable(Throwable throwable) {
        return new WrapResult(throwable);
    }

    private void failHandle(Throwable e) {
        if (this.setDefault) {
            if (completeValue(this.defaultValue)) {
                countDownLatch.countDown();
                tryRunNextWork();
            }
        } else {
            if (completeThrowable(e)) {
                countDownLatch.countDown();
                tryRunNextWork();
            }
        }
    }

    private boolean failValueSet(Throwable e) {
        if (this.setDefault) {
            if (completeValue(this.defaultValue)) {
                countDownLatch.countDown();
                return true;
            }
        } else {
            if (completeThrowable(e)) {
                countDownLatch.countDown();
                return true;
            }
        }
        return false;
    }

    private void doWork() {
        boolean ran;
        T value = null;
        if (this.result.get() != null) {
            tryRunNextWork();
            return;
        }
        try {
            // 等依赖线程结束
            waitForDependAllDone();
            if (this.workable != null) {
                value = this.workable.work();
            }
            ran = true;
        } catch (Throwable e) {
            ran = false;
            failHandle(e);
        }
        if (ran) {
            successHandle(value);
        }
    }


    private void waitForDependAllDone() {
        if (this.dependHasTimeOut) {
            for (WorkNode<?> workNode : this.dependWorkNodes) {
                workNode.getResultWithWait();
            }
        }
    }

    private T getResultWithWait() {
        try {
            if (this.result.get() == null) {
                if (this.timeOut != null) {
                    this.future.get(timeOut, TimeUnit.MILLISECONDS);
                } else {
                    this.future.get();
                }
            }
        } catch (Throwable e) {
            failHandle(e);
        }
        Object r = this.result.get();
        if (r instanceof WrapResult) {
            WrapResult wrapResult = (WrapResult) r;
            if (wrapResult.throwable == null) {
                return null;
            }
            throw new ExceptionWrapper(wrapResult.throwable);
        }
        return (T) r;
    }


    private void successHandle(T value) {
        if (completeValue(value)) {
            countDownLatch.countDown();
            tryRunNextWork();
        }
    }

    public boolean dependIsNotEmpty() {
        return this.dependNotEmpty;
    }

    public boolean dependIsEmpty() {
        return !this.dependNotEmpty;
    }

    private void addNextNode(WorkNode<?> workNode) {
        if (this.nextWorkNodes == null) {
            this.nextWorkNodes = new HashSet<>();
        }
        this.nextWorkNodes.add(workNode);
    }

    @Override
    public String toString() {
        return "WorkNode{" +
                "result=" + result.get() +
                ", setDefault=" + setDefault +
                ", defaultValue=" + defaultValue +
                '}';
    }

    public static class Builder<T> {
        private WorkNode<?>[] waitForList;
        private Workable<T> workable;
        private Long timeOut;

        private boolean defaultValueSet = false;

        private T defaultValue;

        public Builder<T> waitFor(WorkNode<?>... workNodes) {
            this.waitForList = workNodes;
            return this;
        }

        public Builder<T> action(Workable<T> workable) {
            this.workable = workable;
            return this;
        }

        public Builder<T> defaultValue(T defaultValue) {
            this.defaultValueSet = true;
            this.defaultValue = defaultValue;
            return this;
        }

        public Builder<T> timeOut(Long timeOut) {
            this.timeOut = timeOut;
            return this;
        }

        public WorkNode<T> build() {
            WorkNode<T> workNode = new WorkNode<>();
            if (waitForList != null && waitForList.length > 0) {
                int len = waitForList.length;
                Set<WorkNode<?>> dependSet = new HashSet<>();
                for (int i = 0; i < len; i++) {
                    dependSet.add(waitForList[i]);
                    waitForList[i].addNextNode(workNode);
                }
                workNode.dependWorkNodes = dependSet.stream()
                        .sorted(Comparator.comparing(WorkNode::getTimeOut, Comparator.nullsLast(Long::compareTo)))
                        .collect(Collectors.toList());
                workNode.dependHasTimeOut = dependSet.stream().anyMatch(e -> e.getTimeOut() != null);
                workNode.dependNotEmpty = true;
            }
            workNode.setDefault = defaultValueSet;
            workNode.defaultValue = defaultValue;
            workNode.timeOut = timeOut;
            workNode.workable = workable;
            return workNode;
        }
    }

    static class WrapResult {
        private Throwable throwable;

        public WrapResult(Throwable throwable) {
            this.throwable = throwable;
        }
    }
}
